import os
import torch
import torch.optim as optim
from torchvision.io import read_image
from torchvision.transforms import functional as F
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision.transforms import Compose, Resize, Normalize
import random
import torchvision.transforms.functional as TF
from models.common import DetectMultiBackend
from models.yolo import Model
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.torch_utils import select_device
from utils.loss import ComputeLoss
from utils.metrics import ap_per_class
from tqdm import tqdm
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np


class VehicleDataset(Dataset):
    def __init__(self, img_dir, anno_dir, transforms=None):
        """
        初始化数据集
        :param img_dir: 图像目录路径
        :param anno_dir: 标注文件目录路径
        :param transforms: 数据变换操作
        """
        self.img_dir = img_dir
        self.anno_dir = anno_dir
        self.transforms = transforms
        self.img_files = os.listdir(img_dir)

    def __len__(self):
        """
        返回数据集的样本数量
        :return: 样本数量
        """
        return len(self.img_files)

    def __getitem__(self, idx):
        """
        根据索引获取数据集中的一个样本
        :param idx: 样本索引
        :return: 图像和对应的标注信息
        """
        img_path = os.path.join(self.img_dir, self.img_files[idx])
        image = read_image(img_path)
        anno_path = os.path.join(self.anno_dir, self.img_files[idx].replace('.jpg', '.txt'))

        boxes = []
        labels = []
        try:
            with open(anno_path, 'r') as f:
                for line in f.readlines():
                    class_id, x_center, y_center, width, height = line.strip().split()
                    xmin = float(x_center) - float(width) / 2
                    ymin = float(y_center) - float(height) / 2
                    xmax = float(x_center) + float(width) / 2
                    ymax = float(y_center) + float(height) / 2
                    boxes.append([xmin, ymin, xmax, ymax])
                    labels.append(int(class_id))
        except FileNotFoundError:
            print(f"标注文件 {anno_path} 未找到")
            raise

        boxes = torch.as_tensor(boxes, dtype=torch.float32)#若数据集读取的数据是列表形式存储，可以用该函数转张量
        labels = torch.as_tensor(labels, dtype=torch.int64)
        image_id = torch.tensor([idx])#生成一个新张量，把列表、元组、Numpy数组中数组复制转移过去
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((len(boxes),), dtype=torch.int64)#全零张量，指定形状
        
        #数据整合
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target


class RandomHorizontalFlip:
    def __call__(self, image, target):
        """
        以50%的概率对图像和目标框进行水平翻转
        :param image: 输入图像张量
        :param target: 包含目标框等信息的字典
        :return: 翻转后的图像和目标信息
        """
        if random.random() < 0.5:
            height, width = image.shape[-2:]
            image = TF.hflip(image)
            bbox = target["boxes"]
            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]
            target["boxes"] = bbox
        return image, target


def load_and_preprocess_data(img_dir, anno_dir, batch_size):
    """
    加载并预处理数据
    :param img_dir: 图像目录路径
    :param anno_dir: 标注文件目录路径
    :param batch_size: 批次大小
    :return: 训练集、验证集和测试集的数据加载器
    """
    transform = Compose([
        Resize((416, 416)),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    transform_with_aug = Compose([
        Resize((416, 416)),
        RandomHorizontalFlip(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    dataset = VehicleDataset(img_dir, anno_dir, transforms=transform_with_aug)
    train_size = int(0.7 * len(dataset))
    val_size = int(0.2 * len(dataset))
    test_size = len(dataset) - train_size - val_size
    train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, val_loader, test_loader


def build_and_configure_model(yolov5_path):
    """
    构建并配置YOLOv5模型
    :param yolov5_path: YOLOv5模型的根目录路径
    :param weights: 预训练权重文件途径
    :param device: 前面选择的计算设备
    :param dnn: 如果为True,表示使用OpenCV的DNN模块加载模型,此处为False
    :param data: 通常用于指定数据集的配置文件路径,这里是None
    :param fp16: True时采用float16加载模型加载推理,但会损失精度
    """
    device = select_device('')
    weights = yolov5_path + '/yolov5s.pt'
    # 加载预训练模型并获取相关信息
    temp_model = DetectMultiBackend(weights, device=device, dnn=False, data=None, fp16=False)
    imgsz = check_img_size(416, s=temp_model.stride.max())#检查并调整图像尺寸，确保符合模型要求。
    """
    在 YOLOv5 模型中，图像尺寸需要满足一定条件，
    例如是步幅的倍数,check_img_size 函数会对输入的尺寸进行检查和必要的调整，返回一个合适的图像尺寸 imgsz。
    这个尺寸在后续数据预处理（如 Resize 操作）或者模型输入时可能会用到，以保证图像尺寸与模型期望的输入尺寸相匹配。
    """
    state_dict = temp_model.state_dict()#存储模型参数的字典

    # 根据配置文件构建模型并加载预训练权重
    model = Model(yolov5_path + '/models/yolov5s.yaml', ch=3, nc=1)#构建YOLOv5模型。传入yolov5s.yaml配置文件路径，ch=3输入图像通道为3，nc=1表示检测的类别数为1
    model.load_state_dict(state_dict)#加载预训练参数state_dict
    model.to(device)
    model.eval()

    return model, device


def train_model(model, train_loader, device, num_epochs):
    """
    训练模型
    :param model: 要训练的模型
    :param train_loader: 训练集数据加载器
    :param device: 计算设备
    :param num_epochs: 训练轮数
    """
    lr = 0.001
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

    hyp = {'box': 0.05, 'cls': 0.5, 'cls_pw': 1.0, 'obj': 1.0, 'obj_pw': 1.0, 'iou_t': 0.2, 'anchor_t': 4.0,
           'fl_gamma': 0.0, 'hsv_h': 0.015, 'hsv_s': 0.7, 'hsv_v': 0.4, 'degrees': 0.0, 'translate': 0.1,
          'scale': 0.5,'shear': 0.0, 'perspective': 0.0, 'flipud': 0.0, 'fliplr': 0.5,'mosaic': 1.0,
           'mixup': 0.0}
    """
    hyp 是一个字典，包含了一系列超参数。这些超参数在计算损失和数据增强过程中起作用：
    box、cls、obj 等参数通常与损失函数中不同部分（如边界框损失、类别损失、目标置信度损失）的权重相关，用于平衡不同损失项对总损失的贡献。
    iou_t 是交并比(IoU)阈值,用于非极大值抑制（NMS）等操作，判断两个边界框之间的重叠程度。
    hsv_h、hsv_s、hsv_v 用于 HSV 颜色空间的数据增强，控制色调（Hue）、饱和度（Saturation）和明度（Value）的变化范围。
    degrees、translate、scale、shear、perspective、flipud、fliplr、mosaic、mixup 等参数与各种数据增强操作相关，如旋转、平移、缩放、剪切、透视变换、上下翻转、左右翻转、Mosaic 数据增强、Mixup 数据增强等。
    """
    compute_loss = ComputeLoss(model, hyp)

    for epoch in range(num_epochs):
        running_loss = 0.0
        model.train()
        pbar = tqdm(train_loader)#可视化训练进度。tqdm会自动显示训练的进度、预计剩余时间等信息
        for i, (images, targets) in enumerate(pbar):
            images = images.to(device)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]#将键值对k:v 加载到device

            optimizer.zero_grad()
            outputs = model(images)
            loss, loss_items = compute_loss(outputs, targets)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            pbar.set_description(f'Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / (i + 1):.4f}')


def evaluate_model(model, val_loader, device):
    """
    在验证集上评估模型
    :param model: 要评估的模型
    :param val_loader: 验证集数据加载器
    :param device: 计算设备
    :return: mAP、精度和召回率
    """
    model.eval()
    pred_boxes = []
    true_boxes = []

    with torch.no_grad():
        pbar = tqdm(val_loader)
        for i, (images, targets) in enumerate(pbar):
            images = images.to(device)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

            outputs = model(images)
            outputs = non_max_suppression(outputs, conf_thres=0.001, iou_thres=0.6)#非极大值抑制，conf_thres置信度阈值，小于就过滤掉；iou_thres是交并比阈值，用于判断两个框是否属于同一个目标，IoU大于阈值，保留置信度高的

            for pred, target in zip(outputs, targets):#zip把outputs和targets中元素一一对应成一组
                if pred is not None:
                    pred_boxes.append(pred[:, :6])
                true_boxes.append(target['boxes'])

        precision, recall, AP, f1, ap_class = ap_per_class(torch.cat(pred_boxes, 0),
                                                           torch.cat(true_boxes, 0),
                                                           labels=None,
                                                           nc=1)
        """
        调用ap_per_class函数(utils.metrics.py中,在YOLOv5里)计算每个类别的精度(precision,预测成功的比例)、召回率(recall,真实正例中被模型成功预测为正例的比例)、平均精度(AP,不同recall下的精度的加权平均)、F1(precision和recall的调和平均数)、以及每个类别的AP值的类别
        torch.cat(pred_boxes, 0)和torch.cat(true_boxes, 0)分别将对应框列表按照0维拼接成一个张量。
        labels表示无标签
        nc为检测类别=1
        """
        mAP = AP.mean()
        print(f'Validation mAP: {mAP:.4f}, Precision: {precision.mean():.4f}, Recall: {recall.mean():.4f}')

    return mAP, precision.mean(), recall.mean()


def test_model(model, test_loader, device):
    """
    在测试集上测试模型
    :param model: 要测试的模型
    :param test_loader: 测试集数据加载器
    :param device: 计算设备
    :return: mAP、精度和召回率
    """
    model.eval()
    test_pred_boxes = []
    test_true_boxes = []

    with torch.no_grad():
        pbar = tqdm(test_loader)
        for i, (images, targets) in enumerate(pbar):
            images = images.to(device)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

            outputs = model(images)
            outputs = non_max_suppression(outputs, conf_thres=0.001, iou_thres=0.6)

            for pred, target in zip(outputs, targets):
                if pred is not None:
                    test_pred_boxes.append(pred[:, :6])
                test_true_boxes.append(target['boxes'])

        test_precision, test_recall, test_AP, test_f1, test_ap_class = ap_per_class(torch.cat(test_pred_boxes, 0),
                                                                                   torch.cat(test_true_boxes, 0),
                                                                                   labels=None,
                                                                                   nc=1)
        test_mAP = test_AP.mean()
        print(f'Test mAP: {test_mAP:.4f}, Precision: {test_precision.mean():.4f}, Recall: {test_recall.mean():.4f}')

    return test_mAP, test_precision.mean(), test_recall.mean()


def plot_detections(image, pred_boxes, true_boxes, class_names=['vehicle']):
    """
    绘制检测结果
    :param image: 图像张量
    :param pred_boxes: 预测的目标框
    :param true_boxes: 真实的目标框
    :param class_names: 类别名称列表
    """
    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)#批次维度、通道数、高度、宽度
    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
    image = (image * std + mean) * 255#将归一化后的图像还原为原始的像素值范围，乘以 255 是因为归一化前的图像像素值通常在 0 - 255 之间。
    image = image.cpu().numpy().astype(np.uint8)#将图像从 PyTorch 张量转换为 NumPy 数组，并将数据类型转换为无符号 8 位整数（np.uint8），这是 matplotlib 库显示图像所要求的数据类型。cpu()使得张量从gpu转到cpu，因为numpy只能将cpu上tensor转化为numpy数组
    image = np.transpose(image, (1, 2, 0))#将图像的维度从 (C, H, W)（通道数、高度、宽度）转换为 (H, W, C)，因为 matplotlib 要求图像的维度顺序是 (高度, 宽度, 通道数)。

    plt.figure(figsize=(10, 10))#(10,10)英寸
    plt.imshow(image)
    ax = plt.gca()#获取当前坐标轴对象，后续在坐标轴上绘制目标框

    for box in true_boxes:
        """
        绘制真实目标框
        tolist创建列表
        matplotlib.patches创建图形
        ax.add_patch(rect)将矩形补丁添加到当前坐标轴上
        """
        xmin, ymin, xmax, ymax = box.tolist()
        rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, linewidth=1, edgecolor='g', facecolor='none')
        ax.add_patch(rect)

    for box in pred_boxes:#预测目标框及其置信度绘制在图像上
        conf = box[4]#置信度
        xmin, ymin, xmax, ymax = box[:4].tolist()
        rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, linewidth=1, edgecolor='r', facecolor='none')
        ax.text(xmin, ymin, f'{class_names[0]}: {conf:.2f}', color='white', bbox=dict(facecolor='r', alpha=0.5))
        ax.add_patch(rect)

    plt.show()


if __name__ == "__main__":
    img_dir = 'data/images'
    anno_dir = 'data/labels'
    batch_size = 4
    yolov5_path = 'path_to_yolov5'

    train_loader, val_loader, test_loader = load_and_preprocess_data(img_dir, anno_dir, batch_size)
    model, device = build_and_configure_model(yolov5_path)

    num_epochs = 50
    train_model(model, train_loader, device, num_epochs)

    val_mAP, val_precision, val_recall = evaluate_model(model, val_loader, device)#验证集评估模型

    images, targets = next(iter(val_loader))#数据加载器，可迭代对象。经iter()后转化为迭代器，可允许逐个访问集合中的元素。这里意思是取单个出来模拟。
    image = images[0].to(device)
    target = targets[0]
    output = model(image.unsqueeze(0))#image中是类似(C,H,W)，并没有批次信息，所以需要通过unsqueeze(0)添加一个维度，例如变成(1,C,H,W)才能满足原evaluate_model()
    output = non_max_suppression(output, conf_thres=0.001, iou_thres=0.6)[0]

    plot_detections(image, output, target['boxes'])

    test_mAP, test_precision, test_recall = test_model(model, test_loader, device)#测试集评估模型